import soundfile as sf
import math
import wave
import numpy as np

import time 

def shift_audio(wave_data, shift = 0):
    data_len = len(wave_data)
    if(shift > 0):
        tail_data = wave_data[data_len - shift:]
        wave_data[shift:] = wave_data[:data_len - shift]
        wave_data[:shift] = tail_data 
    return wave_data 

def add_noise(wave_data, noise_path = None, out_path = None, SNR = 0, sr=16000):
    if SNR != 0:
        random_values = np.random.rand(len(wave_data)).astype("float32")
        # in case overflow
        wave_data_64 = wave_data.astype("int64")
        Ps = np.sum(wave_data_64 ** 2) / len(wave_data)
        Pn1 = np.sum(random_values ** 2) / len(random_values)
        if(Ps < 0):
            np.save("ps_neg", wave_data)
            print(Ps, SNR, Pn1, Ps / (10 ** (SNR / 10) * Pn1))
        k = math.sqrt(Ps / (10 ** (SNR / 10) * Pn1))
        random_values_we_need = random_values * k
        Pn = np.sum(random_values_we_need ** 2) / len(random_values_we_need)

        snr = 10 * math.log10(Ps / Pn)
        
        wave_data = (wave_data + random_values_we_need + 0.5).astype("int16")

        if out_path is not None:
            with wave.open(out_path, "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)
                wf.setframerate(16000)
                wf.writeframes(wave_data)
                wf.close()
        
    return wave_data

if __name__ == "__main__":
    audio = wave.open("./off.wav")
    signal = audio.readframes(-1)
    signal = np.fromstring(signal, 'int16')
    add_noise(signal, "noise.wav", "off_noise.wav", 0)